from xml.dom.minidom import parse
import xml.dom.minidom
import pandas as pd
import dolphindb as ddb
import numpy as np
import os


class IVAnalysis:
    """
    分析来自qwin的个品种的iv，包括波动率期限结构、两个交易日的波动率对比、偏度数值的波动率计算等、波动率曲面
    """
    def __init__(self, file_path):
        self.base_file_path = file_path
        # 从波动率列表读入的标的物价格
        self.underlying_price_data = {}
        # 全品种iv的临时数据
        self.temp_iv_data = {}
        # 根据目录读取交易日数据
        self.plot_dates = os.listdir(file_path)

    def iv_data_load(self, date):
        """
        加载波动率数据表
        param: date = '230310'
        """
        # 通过xml读取波动率数据
        file_path = os.path.join(self.base_file_path, date, 'TheoryPresettleVolaTmp.xml')
        with open(file_path, 'r', encoding='GBK') as f:
            xml_string = f.read()

        # 压缩空白字符
        xml_string = ''.join(line.strip() for line in xml_string.split('\n'))
        # 解析XML文件
        dom_tree = xml.dom.minidom.parseString(xml_string)
        collection = dom_tree.documentElement
        node_list = collection.getElementsByTagName("Underlying")  # Underlying即标的

        # 读取IV数据
        df_all = pd.DataFrame()  # 存放所有波动率数据
        for node in node_list:
            single_asset_iv_data = {}
            underlying_dict = {}
            for child_node in node.childNodes:
                # 示例：UnderlyingMonth Month="202307" UnderlyingPrice="5424" TheoryPresettleVola="4500:0.1507,.."
                underlying_dict.update(
                    {child_node.getAttribute("Month"): child_node.getAttribute("UnderlyingPrice")})
                single_month_iv_data = {}

                for i in child_node.getAttribute("TheoryPresettleVola").split(','):
                    single_month_iv_data[i.split(':')[0]] = i.split(':')[1]

                single_asset_iv_data.update({child_node.getAttribute("Month"): single_month_iv_data})
            self.temp_iv_data.update({node.getAttribute("UnderlyingCode"): single_asset_iv_data})
            self.underlying_price_data.update({node.getAttribute("UnderlyingCode"): underlying_dict})

        flattened_data = []
        for asset, asset_data in self.temp_iv_data.items():
            for month, month_data in asset_data.items():
                for strike, iv in month_data.items():
                    flattened_data.append({
                        'Asset': asset,
                        'Month': month,
                        'Strike': strike,
                        'IV': iv
                    })
        df = pd.DataFrame(flattened_data)
        df['date'] = date
        df.columns = ['symbol', 'month', 'kprice', 'iv', 'date']
        df['kprice'] = df['kprice'].astype(int)
        df['iv'] = df['iv'].astype(float)
        df['date'] = pd.to_datetime(df['date'])

        # Flatten the 'underlying_price_data' dictionary
        flattened_underlying_data = []
        for asset, asset_data in self.underlying_price_data.items():
            for month, price in asset_data.items():
                flattened_underlying_data.append({
                    'Asset': asset,
                    'Month': month,
                    'UnderlyingPrice': price
                })

        # Convert it to a DataFrame
        df_underlying = pd.DataFrame(flattened_underlying_data)
        df_underlying.columns = ['symbol', 'month', 'UnderlyingPrice']
        df_underlying['UnderlyingPrice'] = df_underlying['UnderlyingPrice'].astype(float)

        # Merge the DataFrames based on 'symbol' and 'month'
        merged_df = pd.merge(df, df_underlying, on=['symbol', 'month'], how='inner')
        # print(merged_df)
        df_all = pd.concat([df_all, merged_df])
        df_all['atm_loc'] = np.log(df_all['kprice']/df_all['UnderlyingPrice'])

        return df_all

    def update_settle_iv(self, data: pd.DataFrame):
        dbName = "dfs://impvResult"
        tbName = "settle_iv"
        s = ddb.session("101.95.132.98", 33147, "admin", "123456")
        appender = ddb.tableAppender(dbName, tbName, s)
        num = appender.append(data)
        return num

    def data_append(self):
        """
        追加所有波动率数据进dolphindb
        :return:
        """
        for i in self.plot_dates:
            self.update_settle_iv(self.iv_data_load(i))
            print('finish%s' % i)


if __name__ == '__main__':
    IV_file_path = os.path.join('c:', os.sep, 'users', '111', 'desktop', 'IV', 'TheoryExportVola')  # os.sep:\\
    iv_analysis = IVAnalysis(IV_file_path)
    # iv_analysis.iv_data_load('20230908')
    iv_analysis.data_append()


